 function [W,W_var,Z,invtausq_w,llik_trace,elbo] = modQTL_association_data(V,X,CIS,PRIOR,sigsq,BURNIN,NEPOCH,W,Z,invtausq_w)
%
% [W,W_var,P,invtausq_w,llik,elbo] = F(V,X,CIS,PRIOR,SIGSQ,BURNIN,NEPOCH,W,Z)
%
% [input]
%
% V          = gene x ind activity matrix
% X          = snp x ind genotype matrix
% CIS        = gene x snp sparse adjacency matrix
% PRIOR      = snp x 1 prior log odds ratio
% SIGSQ      = noise variance
% BURNIN     = burn-in period
% NEPOCH     = Gibbs sampling time
%
% For warm start we can feed initial W and Z; otherwise, W will be
% sampled from N(0,0.1) and Z from snp x 1 zero vector.
%
% [output]
%
% W          = gene x snp association matrix
% W_var      = variance of W with respect to variational distribution
% P          = snp x 1 posterior inclusion probability
% invTauSq_W = gene x NBIBBS sampling results
% llik       = log-likelihood
% elbo       = expected log-likelihood bound
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    TOL = 1e-2;     % seems enough level of precision
    P_TRUNC = 1e-4; % avoid zero probability

    [Ngene, Nind] = size(V);
    [Nsnp, n_temp] = size(X);
    assert( n_temp == Nind );
    [m_temp,n_temp] = size(CIS);
    assert( m_temp == Ngene & n_temp == Nsnp );
    assert( numel(PRIOR) == Nsnp);
    clear *temp;

    NCISTOT = nnz(CIS);
    Ncis_on_gene = full(sum(CIS,2)); % gene x 1
    Ncis_on_snp = full(sum(CIS,1)');    % snp x 1

    PRIOR = double(PRIOR);

    % ================================================================
    % initialization of parameters
    warmstart = false;
    if nargin > 7,
        assert(all(size(W) == size(CIS)));
        assert(numel(Z) == Nsnp);        
        assert(numel(invtausq_w) == Ngene);
        fprintf(2,'Initialized by previous results\r');
        warmstart = true;
    else
        invtausq_w = 0.01*Ncis_on_gene;
        [gene_idx, snp_idx] = find(CIS);
        W = sparse(gene_idx, snp_idx, randn(numel(gene_idx),1)*0.1, Ngene, Nsnp);
        Z = zeros(Nsnp,1);
    end

    snps = find(any(W));
    Vt = double(V');

    % ================================================================
    % precompute
    Xt = double(X');
    clear X;
    VXt = double(V)*Xt;
    clear V;
    Xsq = sum(Xt.^2);

    sigsq = double(sigsq);

    % ================================================================
    % posterior SNP selection
    %
    % log mass M(s) = - (1./sigsq)*[ X(s,:)*[V' - \sum_{t~=s} Z(t)*W(:,t)*X(t,:)]' * W(:,s) - 0.5*||X(s,:)||^2 ||W(:,s)||^2]
    %  = - (1./sigsq) * [ T1 - T2 - 0.5 * Xsq(s) * Wsq(s) ]
    %
    %  T1    = X(s,:) * V'*W(:,s)
    %  T2    = X(s,:) * (ZXtWt - Z(s)*X(s,:)'*W(:,s)') * W(:,s)
    %        = X(s,:) * ZXtWt*W(:,s) - Z(s)X(s,:)*X(s,:)'*W(:,s)'* W(:,s)
    %        = X(s,:) * ZXtWt*W(:,s) - Z(s) * Xsq(s) * Wsq(s)
    %  where
    %  ZXtWt = sum_t Z(t)*X(t,:)'*W(:,t)' * W(:,s)
    %
    % in sum,
    %  = -(1./sigsq)*[X(s,:) * (V'*W(:,s) - ZXtWt*W(:,s)) + (Z(s)-0.5)*Xsq(s)*Wsq(s)]
    %
    % Precomputations
    % ================================================================
    % [a] VtW
    %     - ind x snp
    %     - delta VtW = V(g,:)' * W(g,:)
    %
    % [b] ZXtWt = sum_s Z(s) * X(s,:)' * W(:,s)' where W(~cis(s),s) = 0
    %     - ind x gene
    %     - delta ZXtWt = delta Z(s) * X(s,:)' * W(:,s)'
    %
    % [c] Xsq, snp x 1
    % [d] Wsq, snp x 1
    %
    % ================================================================
    % Estimate W matrix
    %
    %           V(g,:)*X(s,:)' - sum_{t~s} Z(t)*W(g,t)*X(t,:)*X(s,:)'
    % W(g,s) = ---------------------------------------------------
    %           X(s,:)*X(s,:)' + tausq(g)^{-2}
    %
    % numerator = V(cis,:)*X(s,:)' -  sum_t Z(t)*W(cis,t)*X(t,:) * X(s,:)' + Z(s)*W(cis,s)*X(s,:)*X(s,:)'
    %
    %              VXt(cis,s) - ZXtWt(:,cis)'*X(s,:) + Z(s)*W(cis,s)*Xsq(s)
    % W(cis,s)  = ---------------------------------------------------------
    %              Xsq(s) + tausq(cis)^{-2}

    VtW = Vt*W;
    ZXtWt = zeros(Nind,Ngene); % no SNP selected initially
    Wsq = sum(W.^2);

    if warmstart,
        for si = 1:numel(snps),
            s = snps(si);
            cis_genes = CIS(:,s);

            if sum(cis_genes) == 0, continue; end

            ZXtWt(:,cis_genes) = ZXtWt(:,cis_genes) + Z(s) * bsxfun(@times, Xt(:,s), W(cis_genes,s).');

        end
    end

    W_var_conditioned = 0*W;

    llik_trace = NaN(BURNIN+NEPOCH,1);

    for iter = 1:(BURNIN+NEPOCH),

        for si = 1:numel(snps),

            s = snps(si);
            cis_genes = CIS(:,s);

            % if sum(cis_genes) == 0, continue; end

            % sample W(:,s)
            wwold = W(:,s);
            denom_w = Xsq(s) + invtausq_w(cis_genes);
            ww = (VXt(cis_genes,s) - ZXtWt(:,cis_genes)'*Xt(:,s) + Z(s)*W(cis_genes,s)*Xsq(s)) ./ denom_w;
            W_var_conditioned(cis_genes,s) = sigsq ./ denom_w; % given Z(s) = 1

            % ww = VXt(g,s) - ZXtWt(:,g)'*Xt(:,s) + Z(s)*W(g,s)*Xsq(s);
            % ww = ww ./ (Xsq(s) + invtausq_w(g));
            % ww = ww + randn .* sigsq ./ (Xsq(s) + invtausq_w(g));

            delta_ww = ww - full(wwold(cis_genes));
            ZXtWt(:,cis_genes) = ZXtWt(:,cis_genes) + Z(s)* bsxfun(@times, Xt(:,s), delta_ww.');
            VtW(:,s) = Vt(:,cis_genes) * ww;

            W(cis_genes,s) = ww;
            Wsq(s) = sum(ww.^2);


            % sample Z(s)
            zold = Z(s);
            mass = Xt(:,s)'*(VtW(:,s) - ZXtWt(:,cis_genes)*W(cis_genes,s)) + (zold-0.5)*Xsq(s)*Wsq(s);
            mass = mass/sigsq + PRIOR(s);
            mass = mass + 0.5*sum(log(1+Xsq(s)./invtausq_w)) - 0.5*Ncis_on_snp(s)*log(sigsq);

            Z(s) = 1./(1+exp(-mass));
            ZXtWt(:,cis_genes) = ZXtWt(:,cis_genes) + (Z(s)-zold) * bsxfun(@times, Xt(:,s), W(cis_genes,s).');

        end

        % update gene-wise shrinkage
        % compute for genes that have at least one SNP nearby
        if iter > BURNIN,
            subset = find(Ncis_on_gene > 0);
            invtausq_w(subset) = sigsq ./ sum(W(subset,:).^2,2) .* Ncis_on_gene(subset);
            invtausq_w(subset) = arrayfun(@(g) max(0.1/Ncis_on_gene(g), min(10*Ncis_on_gene(g), invtausq_w(g))), subset);
        end

        % log likelihood
        rss = sum(sum((Vt - ZXtWt).^2));
        llik = -0.5*rss/sigsq;
        llik_trace(iter) = llik;

        fprintf(2,'Iter = %03d, LLIK = %.4e, Pr(Z) = %.2f\r',iter,llik,mean(Z));

        if iter > BURNIN + 5,
            if abs(mean(llik_trace((iter-5):(iter-1))) - llik) < TOL,
                fprintf(2,'Converged\r');
                llik_trace = llik_trace(1:iter);
                break;
            end
        end
    end

    W_var = bsxfun(@times, (W_var_conditioned + W.^2)', Z) - bsxfun(@times, W'.^2, Z.^2);
    W_var = W_var';

    % expected log-likelihood bound
    elbo = llik - 0.5*Nind*Ngene*log(sigsq) - (0.5/sigsq)*sum(W_var*Xsq');

    % divergence of prior effect size
    [idx_gene, idx_snp] = find(CIS);
    idx = sub2ind(size(CIS), idx_gene, idx_snp);
    val = log(W_var_conditioned(idx)) + log(invtausq_w(idx_gene)*sigsq);
    log_ratio = sparse(idx_gene, idx_snp, val, Ngene, Nsnp);

    div = double(CIS>0) - bsxfun(@times, W_var_conditioned + W.^2, invtausq_w)/sigsq;
    div = bsxfun(@plus, div, log_ratio);
    elbo = elbo + 0.5*sum(div*Z);

    % divergence of variable selection
    p = Z;
    p(p < P_TRUNC) = P_TRUNC;
    p(p > 1 - P_TRUNC) = 1 - P_TRUNC;
    elbo = elbo + PRIOR'*Z - sum(log(1+exp(PRIOR)));
    elbo = elbo - p'*log(p) - (1-p)'*log(1-p);

end
